package com.mdd.common.plugin.ai;

import com.alibaba.fastjson2.JSONObject;
import com.mdd.common.plugin.ai.request.ChatGptRequest;
import com.mdd.common.plugin.ai.response.ChatGptResponse;
import com.mdd.common.util.RedisUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.util.Assert;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.Proxy;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;

/**
 * @author XYueong
 * @date 2023-12-05
 * @description: TODO
 */
@Slf4j
@Component
public class AiDriver {

    /**
     * 描述: 设置ai
     * @author: XYueong
     * @date: 2023/12/5
     */
    public void  setChatGptApiKey(String key,String chatModel){
        if (key==null){
            throw new IllegalArgumentException("错误 key 不能为空");
        }
        if (chatModel==null){
            throw new IllegalArgumentException("错误 语言模型 不能为空");
        }
        RedisUtils.set("config:openAiApiKey",key);
        RedisUtils.set("config:openAiApiChatModel",chatModel);
    }

    /**
     * 描述: 设置ai 和 代理
     * @author: XYueong
     * @date: 2023/12/7
     */
    public void  setChatGptApiKey(String key,String chatModel,String proxyHost,Integer proxyPort){
        if (key==null){
            throw new IllegalArgumentException("错误 key 不能为空");
        }
        if (chatModel==null){
            throw new IllegalArgumentException("错误 语言模型 不能为空");
        }
        if (proxyHost==null||proxyPort==null){
            throw new IllegalArgumentException("错误 代理地址 填写错误");
        }
        RedisUtils.set("config:openAiApiKey",key);
        RedisUtils.set("config:openAiApiChatModel",chatModel);
        RedisUtils.set("config:proxyHost",proxyHost);
        RedisUtils.set("config:proxyPort",proxyPort);
    }


    /**
     * 描述: 核心方法
     * @author: XYueong
     * @date: 2023/12/5
     */
    public  ChatGptResponse chatGptApi(String questionStr,Proxy proxy ) {

        HttpURLConnection connection = null;
        // ChatGPT API Endpoint
        String apiUrl = "https://api.openai.com/v1/chat/completions";
        // API Key
        String apiKey = (String)RedisUtils.get("config:openAiApiKey");
        Assert.notNull(apiKey, "openai的密钥未配置");

        ChatGptRequest request = new ChatGptRequest();
        List<ChatGptRequest.Messages> question = new ArrayList<>();
        ChatGptRequest.Messages messages = new ChatGptRequest.Messages(questionStr);
        question.add(messages);
        request.setMessages(question);

        // 请求的JSON数据
        String jsonData = JSONObject.toJSONString(request);
        try {
            // 创建URL对象
            URL url = new URL(apiUrl);
            // 打开连接
            if (proxy != null) {
                connection = (HttpURLConnection) url.openConnection(proxy);
            } else {
                connection = (HttpURLConnection) url.openConnection();
            }

            // 设置请求方法为POST
            connection.setRequestMethod("POST");
            // 设置请求头
            connection.setRequestProperty("Content-Type", "application/json");
            connection.setRequestProperty("Authorization", "Bearer " + apiKey);
            // 启用输出流，以便将请求体写入到服务器
            connection.setDoOutput(true);
            // 获取输出流对象
            try (OutputStream outputStream = connection.getOutputStream()) {
                // 将JSON数据写入输出流
                outputStream.write(jsonData.getBytes(StandardCharsets.UTF_8));
                outputStream.flush();
            }

            // 获取响应代码
            int responseCode = connection.getResponseCode();
            if (responseCode == 200) {
                // 读取响应内容
                try (BufferedReader reader = new BufferedReader(new InputStreamReader(connection.getInputStream()))) {
                    StringBuilder response = new StringBuilder();
                    String line;
                    while ((line = reader.readLine()) != null) {
                        response.append(line);
                    }
                    return JSONObject.parseObject(response.toString(), ChatGptResponse.class);
                }
            }
        } catch (Exception e) {
            log.error("请求失败", e);
            throw new RuntimeException("请求失败");
        } finally {
            // 关闭连接
            if (connection != null) {
                connection.disconnect();
            }
        }
        throw new RuntimeException("请求失败");
    }



//    public  void  chatGptApiStream(String questionStr, Integer userId, HttpServletResponse response,ChatBaseOpenAiProxyService proxyService){
//        response.setContentType("text/event-stream");
//        response.setCharacterEncoding("UTF-8");
//        response.setHeader("Cache-Control", "no-cache");
//
//        ServletOutputStream outputStream;
//        try {
//            outputStream = response.getOutputStream();
//        } catch (IOException e) {
//            throw new IllegalArgumentException("错误 获取响应体失败");
//        }
//        proxyService.createStreamChatCompletion(questionStr ,userId+"" ,outputStream );
//    }
}
